import subprocess
import numpy as np
import cv2
import threading

import time
from typing import Callable, Optional, Tuple
from collections import deque


class RTMPFrameStreamer:
    """
    从RTSP帧提取器接收帧并推送到RTMP服务器，支持帧序号跟踪和性能监控
    """

    def __init__(self, rtmp_output_url: str, max_queue_size: int = 30, log_interval: int = 100):
        """
        初始化RTMP帧推流器

        :param rtmp_output_url: RTMP输出流地址
        :param max_queue_size: 帧队列最大容量
        :param log_interval: 日志输出间隔（帧数）
        """
        self.rtmp_output_url = rtmp_output_url
        self.max_queue_size = max_queue_size
        self.log_interval = log_interval
        self.ffmpeg_process: Optional[subprocess.Popen] = None
        self.thread: Optional[threading.Thread] = None
        self.running = False
        self.frame_queue = deque()  # (frame, frame_count)
        self.frame_lock = threading.Lock()
        self.last_frame_count = -1
        self.start_time = 0
        self.processed_frames = 0
        self.dropped_frames = 0

    def start(self) -> bool:
        """启动RTMP推流服务"""
        if self.running:
            print("RTMP推流已在运行中")
            return False

        self._start_ffmpeg()
        self.running = True
        self.start_time = time.time()
        self.processed_frames = 0
        self.dropped_frames = 0
        self.last_frame_count = -1
        self.thread = threading.Thread(target=self._push_frames, daemon=True)
        self.thread.start()
        return True

    def stop(self) -> None:
        """停止RTMP推流服务并输出统计信息"""
        self.running = False
        if self.thread:
            # self.thread.join(timeout=5.0)
            self.thread = None
        if self.ffmpeg_process:
            self.ffmpeg_process.stdin.close()
            self.ffmpeg_process.terminate()
            try:
                self.ffmpeg_process.wait(timeout=1)
            except subprocess.TimeoutExpired:
                self.ffmpeg_process.kill()
            self.ffmpeg_process = None

        # 输出会话统计信息
        elapsed_time = time.time() - self.start_time
        if elapsed_time > 0:
            fps = self.processed_frames / elapsed_time
            drop_rate = (self.dropped_frames / (self.processed_frames + self.dropped_frames)) * 100
            print(f"推流会话统计: 总处理帧={self.processed_frames}, 丢弃帧={self.dropped_frames}, \
                  丢帧率={drop_rate:.2f}%, 平均FPS={fps:.2f}, 持续时间={elapsed_time:.2f}s")
        else:
            print("推流会话已停止")

    def _start_ffmpeg(self) -> None:
        """启动FFmpeg推流进程"""
        ffmpeg_cmd = [
            "ffmpeg",
            "-i", "-",  # 从标准输入读取输入流
            '-f', 'flv',
            "-c:v", "h264_nvenc",  # 使用NVIDIA硬件编码H.264视频
            "-preset", "p2",  # 设置编码预设为p7（最慢但质量最高）
            "-tune", "hq",  # 针对高质量编码进行优化
            "-r", "15",
            "-s", "1280x720",
            #"-s", "1920x1080",
            "-g", "30",  # 关键帧间隔（建议设为2倍帧率，如25fps则设为50）
            "-flush_packets", "1",  # 尽快刷新输出数据包
            "-vsync", "cfr",  # 使用恒定帧率模式
            "-copyts",  # 保留输入时间戳
            "-b:v", "2000k",
            "-maxrate", "4000k",
            "-minrate", "2000k",
            "-bufsize", "8000k",
            self.rtmp_output_url
        ]
        # ffmpeg_cmd = [
        #     "ffmpeg",
        #     "-re",
        #     "-f", "rawvideo",
        #     "-pix_fmt", "bgr24",
        #     "-s", "1920x1080",
        #     "-r", "25",
        #     "-i", "pipe:0",
        #     "-c:v", "h264_nvenc",
        #     "-preset", "p4",
        #     "-tune", "hq",
        #     "-b:v", "4000k",
        #     "-maxrate", "4000k",
        #     "-minrate", "2000k",
        #     "-bufsize", "8000k",
        #     "-f", "flv",
        #     self.rtmp_output_url
        # ]
        # ffmpeg_cmd = [
        #     "ffmpeg",
        #     "-f", "image2pipe",  # 关键：指定输入是图片管道流
        #     "-vcodec", "mjpeg",  # 关键：输入是JPEG编码的图片
        #     "-r", "25",  # 关键：设置与原始RTSP流相同的帧率（需根据实际流调整，比如25/30）
        #     "-copyts",  # 保留输入时间戳
        #     "-i", "-",  # 从标准输入读取
        #     "-c:v", "h264_nvenc",
        #     "-preset", "p2",
        #     "-g", "50",  # 关键帧间隔（建议设为2倍帧率，如25fps则设为50）
        #     "-b:v", "2000k",
        #     "-maxrate", "4000k",
        #     "-minrate", "2000k",
        #     "-bufsize", "8000k",
        #     "-f", "flv",
        #     self.rtmp_output_url
        # ]

        self.ffmpeg_process = subprocess.Popen(
            ffmpeg_cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )

    def _push_frames(self) -> None:
        """从队列获取帧并推送到RTMP服务器"""
        while self.running:
            with self.frame_lock:
                if self.frame_queue:
                    frame, frame_count = self.frame_queue.popleft()
                else:
                    frame = None

            if frame is not None:
                try:
                    _, buffer = cv2.imencode('.jpg', frame)
                    # self.ffmpeg_process.stdin.write(f"frame_id={frame_count}\n".encode() + buffer.tobytes())
                    self.ffmpeg_process.stdin.write(buffer.tobytes())
                    self.ffmpeg_process.stdin.flush()
                    self.processed_frames += 1

                    # 定期输出性能日志
                    if self.processed_frames % self.log_interval == 0:
                        elapsed = time.time() - self.start_time
                        fps = self.processed_frames / elapsed
                        print(f"[性能监控] 已处理 {self.processed_frames} 帧, FPS={fps:.2f}, 队列剩余={len(self.frame_queue)}")

                except Exception as e:
                    print(f"推流错误: {e}")
                    self.stop()
                    self.start()
                    break
            else:
                threading.Event().wait(0.01)

    def handle_frame(self, frame: np.ndarray, frame_count: int) -> None:
        """
        帧处理回调函数，由RTSPFrameExtractorByFFmpeg调用

        :param frame: 接收到的视频帧
        :param frame_count: 帧序号
        """
        if not self.running:
            return

        # 检查帧序号连续性，检测丢帧
        if self.last_frame_count != -1 and frame_count != self.last_frame_count + 1:
            lost_frames = frame_count - self.last_frame_count - 1
            print(f"[WARNING] 检测到丢帧: 丢失 {lost_frames} 帧 (预期 {self.last_frame_count + 1}, 实际 {frame_count})")
        self.last_frame_count = frame_count

        with self.frame_lock:
            # 队列满时丢弃最旧帧保留新帧
            if len(self.frame_queue) >= self.max_queue_size:
                self.frame_queue.popleft()
                self.dropped_frames += 1
            self.frame_queue.append((frame, frame_count))

